From 3117695a07a5c586d079fccb821628a7a0710653 Mon Sep 17 00:00:00 2001
From: "Karl O. Pinc" <kop@karlpinc.com>
Date: Sun, 29 Sep 2024 17:55:42 -0500
Subject: [PATCH] Rewrite so downloads do not put sql results in RAM

---
 src/pgwui_sql/views/sql.py | 77 ++++++++++++++++++++++----------------
 1 file changed, 44 insertions(+), 33 deletions(-)

diff --git a/src/pgwui_sql/views/sql.py b/src/pgwui_sql/views/sql.py
index 9c1d077..c39fc34 100644
--- a/src/pgwui_sql/views/sql.py
+++ b/src/pgwui_sql/views/sql.py
@@ -247,10 +247,8 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
                 detail=self.format_detail(err, stmt_text))
 
     def get_result_rows(self, cur, null_rep):
-        rows = []
         while (row := cur.fetchone()) is not None:
-            rows.append(ResultRow().build_data_row(row, null_rep))
-        return rows
+            yield ResultRow().build_data_row(row, null_rep)
 
     def get_db_search_path(self):
         self.cur.execute('SHOW search_path;')
@@ -277,20 +275,43 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
             quoting = csv.QUOTE_NONNUMERIC
         return csv.writer(self.tfile, quoting=quoting)
 
-    def make_download(self, sql_results):
-        downloading = self.uf['download']
-        if downloading:
-            self.tfile = tempfile.TemporaryFile(mode='w+t', newline='')
-            writer = self.make_csv_writer()
-            if self.uf['include_sql']:
-                writer.writerow((self.uf['sql'].rstrip(),))
-            for sql_result in sql_results:
-                writer.writerow((sql_result.statusmessage.data,))
-                writer.writerow(sql_result.heading.data)
-                for row in sql_result.rows:
+    def make_download(self, cur):
+        # Optimized to minimize RAM usage
+        null_rep = self.uf['null_rep']
+        self.tfile = tempfile.TemporaryFile(mode='w+t', newline='')
+        writer = self.make_csv_writer()
+        if self.uf['include_sql']:
+            writer.writerow((self.uf['sql'].rstrip(),))
+
+        nextset = True
+        while nextset is True:
+            # Rather than report the statusmessage first, which requires
+            # putting all the statement's results in RAM, report it last.
+            if cur.rownumber is not None:
+                writer.writerow(ResultRow().build_heading_row(cur).data)
+                for row in self.get_result_rows(cur, null_rep):
                     writer.writerow(row.data)
-                writer.writerow((sql_result.rowcount.data,))
-            self.tfile.seek(0)
+            writer.writerow((ResultRow().build_rowcount_row(cur).data,))
+            writer.writerow((ResultRow().build_statusmessage_row(cur).data,))
+
+            nextset = cur.nextset()
+
+        self.tfile.seek(0)
+
+    def make_sql_results(self, cur):
+        null_rep = self.uf['null_rep']
+
+        sql_results = self.sql_results
+        nextset = True
+        while nextset is True:
+            sql_result = SQLResult()
+            if cur.rownumber is not None:
+                sql_result.heading = ResultRow().build_heading_row(cur)
+                sql_result.rows = list(self.get_result_rows(cur, null_rep))
+            sql_result.statusmessage = ResultRow().build_statusmessage_row(cur)
+            sql_result.rowcount = ResultRow().build_rowcount_row(cur)
+            sql_results.append(sql_result)
+            nextset = cur.nextset()
 
     def cleanup(self):
         '''
@@ -300,7 +321,6 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
         should there be errors the results that do exist are displayed.
         '''
         cur = self.cur
-        null_rep = self.uf['null_rep']
 
         # Adjust the executed SQL to use the requested search_path
         # Change the form content so that the user sees the change
@@ -313,22 +333,13 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
         if cur.statusmessage is None:
             raise sql_ex.NoStatementsError(descr='No SQL statements executed')
 
-        sql_results = self.sql_results
-        nextset = True
-        while nextset is True:
-            sql_result = SQLResult()
-            if cur.rownumber is not None:
-                sql_result.heading = ResultRow().build_heading_row(cur)
-                sql_result.rows = self.get_result_rows(cur, null_rep)
-            sql_result.statusmessage = ResultRow().build_statusmessage_row(cur)
-            sql_result.rowcount = ResultRow().build_rowcount_row(cur)
-            sql_results.append(sql_result)
-            nextset = cur.nextset()
-
-        try:
-            self.make_download(sql_results)
-        except csv.Error as ex:
-            raise sql_ex.CSVError(descr=f'The csv module reports: {ex}')
+        if self.uf['download']:
+            try:
+                self.make_download(cur)
+            except csv.Error as ex:
+                raise sql_ex.CSVError(descr=f'The csv module reports: {ex}')
+        else:
+            self.make_sql_results(cur)
 
     def factory(self, ue):
         '''Make a db loader function from an UploadEngine.
-- 
2.34.1